feat(regularization): numba kernels for the numpy split-regularization assembly (#267 lever 1) - #553

Summary

Lever 1 of PyAutoLabs/autolens_profiling#267 (phase 3 of the fixed-lens-light-numba-cpu epic): the numpy/numba split-regularization assembly is now two @numba_util.jit() kernels instead of two pure-Python loops.

On the production numba CPU path at HST Delaunay N=1500, inversion.regularization_matrix was the single largest site in the source-only likelihood call: 111.5 ms of 404.6 ms (27.6 %), invariant under the NNLS memo. regularization_util.py contained zero numba; PyAutoArray#536 optimised the JAX branch only.

What changed

• reg_split_np_from and pixel_splitted_regularization_matrix_np_from keep their names and signatures and now call _reg_split_kernel / _pixel_splitted_regularization_matrix_kernel, which reproduce the original accumulation order bit-for-bit (the 2e-8 diagonal first, both-triangle adds in l / l+m order, the trailing /= 2.0 last; fastmath=False).
• The original Python bodies are retained verbatim as _reg_split_reference and _pixel_splitted_regularization_matrix_reference. They remain the ground truth for the JAX branch tests (test_pixel_splitted_jax.py) and are now also pinned against the kernels with np.array_equal in the new test_regularization_util_numba.py (8 tests: narrow/wide DelaunayNN-shaped tables, K-NN tables through reg_split, no input mutation, the size == 0 row, the overflow raise, and an end-to-end ConstantSplit.regularization_matrix_from against the reference chain).
• The JAX branch of pixel_splitted_regularization_matrix_from / reg_split_from is byte-identical to main.

API Changes

Three deliberate behaviour changes on the numpy path, none of which affects any caller in PyAutoArray, PyAutoGalaxy, PyAutoLens, PyAutoCTI or PyAutoFit (all consume the returned tuple):

1. reg_split_np_from no longer mutates its arguments. It copies splitted_mappings and splitted_sizes on entry and returns the completed tables; the interpolators' cached_property stencil tables are no longer rewritten by the first regularization-matrix call. The returned mappings keep their full width including the reserved pad column, so the (mappings, sizes, weights) contract is unchanged. Side effect: test_adapt_power_jax.py now compares its numpy and JAX legs from the same inputs (previously the JAX leg ran on already-mutated tables).
2. An empty stencil row (size == 0) inserts its centre pixel at column 0. The old loop left j at the previous row's value and wrote at a stale j + 1. The kernel uses an explicit insert = sizes[i].
3. A full stencil with no self pixel raises exc.InversionException naming the row and table width, instead of an IndexError (Python) or silent out-of-bounds write (numba has no bounds checking).

Measurements

Local single-thread sanity (P=1500, K=4, warm, median of 20): reference 106.3 ms → kernel 1.23 ms (86x); reg_split 9.13 → 0.040 ms.

RAL A/B (gpu partition CPUs-only, 1 thread numba + BLAS, private merge-base control 5e2bc0f4 vs feature on PYTHONPATH, phase-2 Leg D cell: route b, memo ON):

1.38x on the whole call (1.383x on medians); the H site falls from 26.8 % to 2.0 % of the call. Both arms ran on one node back-to-back; the feature arm entered at a higher load average than the control, so 1.38x is a lower bound. The control arm reproduces phase 2's Leg D (404.6 ms / 111.5 ms) to +2.2 % / −0.8 %.

Witness inside the feature arm on RAL (fixed_light_numba_levers_witness.py): the library's inversion.regularization_matrix is array_equal to the retained reference chain (max |diff| 0.0); the interpolator's cached stencil tables are byte-identical after the inversion; the fit's log evidence with the kernels vs the monkeypatched references differs by exactly 0.0 (gate ≤ 1e-9).

A100 row (job 343346, delaunay.py breakdown cell, fp64, AdaptSplit): identity, as expected. The GPU cell dispatches to the JAX branch, so the numpy kernels are never reached there. The H prefix agrees to 0.03 % (7.3318 vs 7.3298 ms) and the whole call to 0.04 %; log_evidence_terms are bit-identical between arms. Finding for the record: with InterpolatorDelaunay the stencil table is K = 4, so SPLIT_REG_COMPACT_WIDTH = 12 / SPLIT_REG_WIDE_ROW_BUDGET = 256 are inert at HST Delaunay N=1500 and can only be tuned on a cell with K > 12.

Full tables, gate roll-up and provenance land in the paired autolens_profiling PR (results/notes/).

Testing

• pytest test_autoarray/inversion/regularizations -q — 118 passed
• pytest test_autoarray/inversion/pixelization/interpolator -q — 58 passed
• pytest test_autoarray/inversion/inversion -q -x — 212 passed
• black --check clean

Closes nothing on its own; tracked by PyAutoLabs/autolens_profiling#267 (lever 1 of 3).

Heart gate

pyauto-heart readiness was RED (score 20) at ship time with two release-chain reasons unrelated to this branch — "install verification FAILED (testpypi; checks F)" and "release validation FAILED (stage integrate)" — acknowledged by the human on 2026-09-16 for PR-open only. Merge stays a human act via /prm.

🤖 Generated with Claude Code